package com.meatball.game.asr

import com.fasterxml.jackson.databind.DeserializationFeature
import com.fasterxml.jackson.databind.ObjectMapper
import org.java_websocket.client.WebSocketClient
import org.java_websocket.handshake.ServerHandshake
import org.slf4j.LoggerFactory
import java.io.ByteArrayInputStream
import java.io.ByteArrayOutputStream
import java.io.IOException
import java.net.URI
import java.net.URISyntaxException
import java.nio.ByteBuffer
import java.nio.charset.StandardCharsets
import java.util.*
import java.util.concurrent.CountDownLatch
import java.util.concurrent.TimeUnit
import java.util.zip.GZIPInputStream
import java.util.zip.GZIPOutputStream
import javax.crypto.Mac
import javax.crypto.spec.SecretKeySpec

class AsrClient(uri: URI) : WebSocketClient(uri)  {
    companion object {
        private const val URL = "wss://openspeech.bytedance.com/api/v2/asr"
        private val logger = LoggerFactory.getLogger(WebSocketClient::class.java)

        @Throws(URISyntaxException::class)
        fun build(): AsrClient {
            val uri = URI(URL)
            return AsrClient(uri)
        }
    }

    private var appid: String = ""
    private var token: String = ""
    private var sk: String = ""
    private var cluster: String = ""
    private var workflow: String = "audio_in,resample,partition,vad,fe,decode"
    private var uid: String = "user_id"
    private var nhest: Int = 1
    private var show_utterances: Boolean = true
    private var result_type: String = "full"
    private var format: String = "wav"
    private var codec: String = "raw"
    private var sample_rate: Int = 16000
    private var channels: Int = 1
    private var bits: Int = 16
    private var authType: AuthType = AuthType.TOKEN
    private var params_msg: ByteArray = ByteArray(0)
    private var asr_response: AsrResponse? = null
    private var recv_latch = CountDownLatch(0)
    private var recv_timeout: Int = 5
    private var recv_suc: Boolean = true

    enum class AuthType {
        TOKEN,
        SIGNATURE
    }

    override fun onOpen(serverHandshake: ServerHandshake) {
        logger.info("asr client onOpen")
    }

    override fun onMessage(s: String) {
        logger.info("onMessage String, should be onMessage(ByteBuffer) called")
    }

    override fun onMessage(bytes: ByteBuffer) {
        try {
            if (parseResponse(bytes) != 0) {
                recv_suc = false
                logger.error("error happens to close connection")
                close()
            }
            recv_latch.countDown()
        } catch (e: IOException) {
            e.printStackTrace()
        }
    }

    override fun onClose(i: Int, s: String, b: Boolean) {
        logger.info("asr onClose {}, {}, {}", i, s, b)
    }

    override fun onError(e: Exception) {
        logger.info("asr onError {}", e.message)
        recv_suc = false
        recv_latch.countDown()
        close()
    }

    fun asrSyncConnect(): Boolean {
        params_msg = constructParam()
        setAuthHeader()
        var ret = connectBlocking()
        if (!ret) {
            return ret
        }
        recv_latch = CountDownLatch(1)
        send(params_msg)
        ret = recv_latch.await(recv_timeout.toLong(), TimeUnit.SECONDS)
        return ret == true && recv_suc
    }

    fun asrSend(audio: ByteArray, isLast: Boolean): AsrResponse {
        recv_latch = CountDownLatch(1)
        val payload = constructAudioPayload(audio, isLast)
        send(payload)
        val ret = recv_latch.await(recv_timeout.toLong(), TimeUnit.SECONDS)
        if (ret != true) {
            logger.error("recv message timeout")
            close()
            return AsrResponse()
        }
        return asr_response ?: AsrResponse()
    }

    fun asrClose(): Int {
        close()
        return 0
    }

    private fun setAuthHeader() {
        if (authType == AuthType.TOKEN) {
            addHeader("Authorization", "Bearer; $token")
            return
        }

        val custom_header = "Custom"
        val custom_cont = "auth_custom"
        addHeader(custom_header, custom_cont)

        val str = "GET ${uri.path} HTTP/1.1\n$custom_cont\n"
        val str_byte = str.toByteArray(StandardCharsets.UTF_8)
        val data = concatByte(str_byte, params_msg)

        val sk_byte = sk.toByteArray(StandardCharsets.UTF_8)
        val HMAC_SHA256 = "HmacSHA256"
        val sha256Hmac = Mac.getInstance(HMAC_SHA256)
        val keySpec = SecretKeySpec(sk_byte, HMAC_SHA256)
        sha256Hmac.init(keySpec)
        val mac_data = sha256Hmac.doFinal(data)

        val base64_data = Base64.getUrlEncoder().encodeToString(mac_data)
        val auth_cont = "HMAC256; access_token=\"$token\"; mac=\"$base64_data\"; h=\"$custom_header\""
        addHeader("Authorization", auth_cont)
    }

    private fun gzipCompress(content: ByteArray): ByteArray {
        val out = ByteArrayOutputStream(content.size)
        val gzip = GZIPOutputStream(out)
        gzip.write(content)
        gzip.close()
        val result = out.toByteArray()
        out.close()
        return result
    }

    private fun gzipDecompress(content: ByteArray): ByteArray {
        val `in` = ByteArrayInputStream(content)
        val gzip = GZIPInputStream(`in`)
        val out = ByteArrayOutputStream()
        val buff = ByteArray(1024)
        var len: Int
        while (gzip.read(buff, 0, buff.size).also { len = it } > 0) {
            out.write(buff, 0, len)
        }
        val result = out.toByteArray()
        `in`.close()
        gzip.close()
        out.close()
        return result
    }

    private fun constructParam(): ByteArray {
        val header_len = 4
        val header = ByteArray(header_len)
        header[0] = (ProtocolVersion.PROTOCOL_VERSION shl 4 or (header_len shr 2)).toByte()
        header[1] = (MessageType.FULL_CLIENT_REQUEST shl 4 or MessageTypeFlag.NO_SEQUENCE_NUMBER).toByte()
        header[2] = (MessageSerial.JSON shl 4 or MessageCompress.GZIP).toByte()
        header[3] = 0

        val reqid = UUID.randomUUID().toString()
        val app = AsrParams.App(appid, cluster, token)
        val user = AsrParams.User(uid)
        val request = AsrParams.Request(reqid, workflow, 1, show_utterances, result_type, 1)
        val audio = AsrParams.Audio(format, codec, sample_rate, bits, channels)
        val asr_params = AsrParams(app, user, request, audio)
        val mapper = ObjectMapper()
        val payload = mapper.writeValueAsBytes(asr_params)
        payload?.let {
            logger.info("params_json {}", String(it))
            val compressedPayload = gzipCompress(it)

            val payload_len = compressedPayload.size
            val bb = ByteBuffer.allocate(4)
            bb.putInt(payload_len)
            val pl_byte = bb.array()

            return concatByte(header, pl_byte, compressedPayload)
        }
        return ByteArray(0)
    }

    private fun parseResponse(msg: ByteBuffer): Int {
        val msg_byte = msg.array()
        val header_len = (msg_byte[0].toInt() and 0x0f) shl 2
        val message_type = (msg_byte[1].toInt() and 0xf0) shr 4
        val message_type_flag = msg_byte[1].toInt() and 0x0f
        val message_serial = (msg_byte[2].toInt() and 0xf0) shr 4
        val message_compress = msg_byte[2].toInt() and 0x0f
        var payload: ByteArray? = null
        var payload_len = 0
        var payload_offset = header_len

        if (message_type == MessageType.FULL_SERVER_RESPONSE) {
            val bb = ByteBuffer.wrap(msg_byte, payload_offset, 4)
            payload_len = bb.int
            payload_offset += 4
        } else if (message_type == MessageType.SERVER_ACK) {
            val bb = ByteBuffer.wrap(msg_byte, payload_offset, 4)
            val seq = bb.int
            payload_offset += 4
            if (msg_byte.size > 8) {
                payload_len = ByteBuffer.wrap(msg_byte, payload_offset, 4).int
                payload_offset += 4
            }
        } else if (message_type == MessageType.ERROR_MESSAGE_FROM_SERVER) {
            val error_code = ByteBuffer.wrap(msg_byte, payload_offset, 4).int
            payload_offset += 4
            payload_len = ByteBuffer.wrap(msg_byte, payload_offset, 4).int
            payload_offset += 4
        } else {
            logger.error("unsupported message type {}", message_type)
            return -1
        }

        payload = ByteArray(msg_byte.size - payload_offset)
        System.arraycopy(msg_byte, payload_offset, payload, 0, payload.size)
        if (message_compress == MessageCompress.GZIP) {
            payload = gzipDecompress(payload)
        }
        if (message_serial == MessageSerial.JSON) {
            val mapper = ObjectMapper().disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES)
            asr_response = mapper.readValue(payload, AsrResponse::class.java)
        }
        if (asr_response?.code != 1000) {
            logger.error("asr response {}", String(payload))
            return -1
        }
        if (asr_response?.sequence ?: 0 < 0) {
            logger.debug("get last response")
        }
        logger.info("asr response {}", String(payload))

        return 0
    }

    private fun constructAudioPayload(audio: ByteArray, isLast: Boolean): ByteArray {
        val header_len = 4
        val header = ByteArray(header_len)
        header[0] = (ProtocolVersion.PROTOCOL_VERSION shl 4 or (header_len shr 2)).toByte()
        header[1] = if (!isLast) {
            (MessageType.AUDIO_ONLY_CLIENT_REQUEST shl 4 or MessageTypeFlag.NO_SEQUENCE_NUMBER).toByte()
        } else {
            (MessageType.AUDIO_ONLY_CLIENT_REQUEST shl 4 or MessageTypeFlag.NEGATIVE_SEQUENCE_SERVER_ASSGIN).toByte()
        }
        header[2] = (MessageSerial.JSON shl 4 or MessageCompress.GZIP).toByte()
        header[3] = 0

        val payload = gzipCompress(audio)
        val payload_len = payload.size
        val bb = ByteBuffer.allocate(4)
        bb.putInt(payload_len)
        val pl_byte = bb.array()

        return concatByte(header, pl_byte, payload)
    }

    private fun concatByte(first: ByteArray, second: ByteArray): ByteArray {
        val result = ByteArray(first.size + second.size)
        System.arraycopy(first, 0, result, 0, first.size)
        System.arraycopy(second, 0, result, first.size, second.size)
        return result
    }

    private fun concatByte(first: ByteArray, second: ByteArray, third: ByteArray): ByteArray {
        val result = ByteArray(first.size + second.size + third.size)
        System.arraycopy(first, 0, result, 0, first.size)
        System.arraycopy(second, 0, result, first.size, second.size)
        System.arraycopy(third, 0, result, first.size + second.size, third.size)
        return result
    }

    fun setAppid(appid: String) {
        this.appid = appid
    }

    fun setToken(token: String) {
        this.token = token
    }

    fun setSk(sk: String) {
        this.sk = sk
    }

    fun setCluster(cluster: String) {
        this.cluster = cluster
    }

    fun setWorkflow(workflow: String) {
        this.workflow = workflow
    }

    fun setUid(uid: String) {
        this.uid = uid
    }

    fun setShowUtterances(show_utterances: Boolean) {
        this.show_utterances = show_utterances
    }

    fun setResultType(result_type: String) {
        this.result_type = result_type
    }

    fun setFormat(format: String) {
        this.format = format
    }

    fun setCodec(codec: String) {
        this.codec = codec
    }

    fun setSampleRate(sample_rate: Int) {
        this.sample_rate = sample_rate
    }

    fun setChannels(channels: Int) {
        this.channels = channels
    }

    fun setBits(bits: Int) {
        this.bits = bits
    }

    fun setAuthType(authType: AuthType) {
        this.authType = authType
    }

    fun getAsrResponse(): AsrResponse? {
        return asr_response
    }
}